#!/usr/bin/env python

import glob
import os
import optparse
import random
import sys

import numpy
import tables

import charmicat
import histogram

class FillWorker(charmicat.MPWorker):
    def work(self, colc, hcoll, paths, opts):                
        bufsize = max([h._bufsize for h in hcoll])
        for h in hcoll:
            h.set_buffer_size(bufsize)

        self.nfile = len(paths)

        assoc_cols = hcoll.assoc_cols()

        for ip, path in enumerate(paths):
            fpin = tables.openFile(path)
            cols_in_data = set(fpin.root.nt.colnames) & assoc_cols
            cols_to_produce = assoc_cols - cols_in_data
            if len(cols_to_produce) > 0:
                cols_to_load = cols_in_data | colc.v4cols
            else:
                cols_to_load = cols_in_data
            
            self.log(ip)
            nrows = len(fpin.root.nt)

            for i in xrange(0, nrows, bufsize):
                stop = i + bufsize 
                if (i + bufsize) > nrows:
                    stop = nrows
                    
                coldict = {}
                for cid in cols_to_load:
                    coldict[cid] = numpy.nan_to_num(fpin.root.nt.read(start=i, 
                                                                     stop=stop, 
                                                                     field=cid))
                if len(cols_to_produce) > 0:
                    vecs = colc.four_vectors(coldict)

                    for c2p in cols_to_produce:
                        coldict[c2p] = colc[c2p](coldict, vecs)
                    
                hcoll.fill(coldict)
                            
            fpin.close()
        
        self.log(len(paths))

        outpath = os.path.join(self.logdir, str(os.getpid())) + '.hpt'
        histogram.dump(hcoll, outpath)

        
    

def main():
    """
    ccat_module_hists.py [opts] file1.h5 [file2.h5 [...]]

    Fill histograms from a module supplied from the command line or inferred 
    from the file.
    """

    op = optparse.OptionParser(usage=main.__doc__)


    op.add_option("-c", "--collection", dest='hcolls', action='append', 
                  default=[], metavar="COLLECTION",
                  help="Fill histogram collection COLLECTION" \
                      + " (can be repeated)")                  
    op.add_option("-l", "--logdir", dest='logdir', default=None,
                  help='Specify directory PATH for MP logging', 
                  metavar='PATH')
    op.add_option("-m", "--module", dest='module_name', default=None, 
                  help="Use module MODULE instead of inferring from input", 
                  metavar='MODULE')
    op.add_option("-n", "--name", dest="list", action='append', default=[],
                  help="Fill histogram NAME (can be used multiple times, " \
                      + "compatible w/ --histset",
                  metavar="NAME")
    op.add_option("-o", "--output", dest="output", default=None, 
                  help="Place output in FILE.hpt", metavar="FILE")
    op.add_option("-v", "--verbose", dest='verbose', default=False, 
                  action='store_true', help='Print more output')
    op.add_option("--cpus", dest="ncpu", type="int", 
                  default=charmicat.DEFAULT_NCPU,
                  help="Spawn no more than N threads", metavar="N")
    op.add_option("--force-log", dest='force_log', default=False, 
                  action='store_true',
                  help='Force reuse of currently-existing log directory')

    opts, args = op.parse_args()

    if len(args) == 0:
        msg = "Please supply at least one input file."
        raise charmicat.BadCommandLineError(msg, -1)

    if not opts.module_name:
        opts.module_name = charmicat.module_name_from_paths(args)
    
    try:
        module = __import__(opts.module_name)           
    except ImportError, e:
        msg = "Error importing module `{0}': \n{1}".format(opts.module_name, e)
        raise charmicat.BadCommandLineError(msg, -2)

    hcoll = charmicat.build_histcoll(module.histdef, opts)

    if opts.output is None:
        opts.output = os.path.basename(os.getcwd())

    if len(args) < opts.ncpu:
        opts.ncpu = len(args)
    
    opts.output = os.path.abspath(opts.output)

    table_columns = charmicat.column_names_from_paths(args)    
    column_fillable = hcoll.column_fillable(table_columns)
    vector_fillable = hcoll.not_column_fillable(table_columns)

    if opts.verbose:
        print "Column-fillable histograms: "
        charmicat.column_output([h.name for h in column_fillable], 
                                prefix='\t')
        print "Vector-filled histograms: "
        charmicat.column_output([h.name for h in vector_fillable], 
                                prefix='\t')

    charmicat.assert_fillable(module.column_calc, hcoll, table_columns)
                              
    input_lists = charmicat.n_groups(opts.ncpu, args)

    arg_lists = []
    for il in input_lists:
        arg_lists.append((module.column_calc, hcoll, il, opts))

    monitor = charmicat.FillMonitor(opts.ncpu, len(args), opts.logdir, opts)
    workers = [FillWorker(monitor.logdir, args=al) for al in arg_lists]

    charmicat.mp_deploy_and_monitor(monitor, workers, ncpu=opts.ncpu)

    monitor.hcoll.module_name = opts.module_name
    histogram.dump(monitor.hcoll, opts.output)

    return
        


if __name__ == "__main__":
    try:
        sys.exit(main())
    except charmicat.BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)

